#include "net.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

static struct ibv_qp *create_qp_init(struct ibv_pd *pd, struct ibv_cq *recv_cq, 
                            struct ibv_cq *send_cq, int ib_port){
  struct ibv_qp_init_attr c_attr;
  struct ibv_qp *qp;
  struct ibv_qp_attr i_attr;
  int res;

  memset(&c_attr, 0, sizeof c_attr);
  c_attr.send_cq = send_cq;
  c_attr.recv_cq = recv_cq;
  c_attr.qp_type = IBV_QPT_RC;
  c_attr.cap.max_send_wr = CQ_DEPTH;
  c_attr.cap.max_recv_wr = CQ_DEPTH;
  c_attr.cap.max_send_sge = MAX_SGE_SIZE;
  c_attr.cap.max_recv_sge = MAX_SGE_SIZE;
  c_attr.cap.max_inline_data = 0; 
  c_attr.sq_sig_all = 0;
  qp = ibv_create_qp(pd, &c_attr);
  assert(qp);
  memset(&i_attr,0,sizeof(struct ibv_qp_attr));
  i_attr.qp_state = IBV_QPS_INIT;
  i_attr.port_num = ib_port;
  i_attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | 
                                IBV_ACCESS_REMOTE_WRITE |
                                IBV_ACCESS_REMOTE_READ |
                                IBV_ACCESS_REMOTE_ATOMIC;
  res = ibv_modify_qp(qp, &i_attr ,IBV_QP_STATE | 
                          IBV_QP_PKEY_INDEX | IBV_QP_PORT |
                          IBV_QP_ACCESS_FLAGS);
  assert(res == 0);
  return qp;
}


void get_qp_info(struct mc_ctx *mi, const char *prefix, int rmt_id, int loc_id, struct qp_info *info){
  struct qp_info *qi;
  char name[SHORT_STR_LEN];
  snprintf(name, SHORT_STR_LEN, "%s-%d-%d", prefix, rmt_id, loc_id);
  if(strlen(name) >= SHORT_STR_LEN - 1)p_die("prefix too long\n");
  p_info("get qp : %s\n", name);
  while(mt_get(mi, name, (char **)&qi, NULL))usleep(1000);
  memcpy(info, qi, sizeof(struct qp_info));
  free(qi);
}

void connect_qp(int ib_port ,struct ibv_qp *qp, struct qp_info *info){
  struct ibv_qp_attr attr;
  int res;

  attr.qp_state = IBV_QPS_RTR;
  attr.path_mtu = IBV_MTU_4096;
  attr.dest_qp_num = info->qpn;
  attr.rq_psn = PSN;
  attr.max_dest_rd_atomic = 16;  // relate to IBV_WC_REM_INV_REQ_ERR
  attr.min_rnr_timer = 12;
  attr.ah_attr.is_global = 0;
  attr.ah_attr.dlid = info->lid;
  attr.ah_attr.sl = 0;
  attr.ah_attr.src_path_bits = 0;
  attr.ah_attr.port_num = ib_port;
  res = ibv_modify_qp(qp, &attr,
              IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU |
              IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
              IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER
              );
  assert(res == 0);
  attr.qp_state = IBV_QPS_RTS;
  attr.timeout = 14;  
  attr.retry_cnt = 7;
  attr.rnr_retry = 7;
  attr.sq_psn = PSN;
  attr.max_rd_atomic = 16; 
  res = ibv_modify_qp(qp, &attr,
              IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
               IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN |
              IBV_QP_MAX_QP_RD_ATOMIC);
  assert(res == 0);
}
static void publish_qp_info(struct server_ctx *ctx ,const char *prefix, 
                      struct ibv_qp *qp, int loc_id, int rmt_id){
  struct qp_info info;
  char name[SHORT_STR_LEN];
  info.lid = ctx->port_attr.lid;
  info.qpn = qp->qp_num;
  snprintf(name, SHORT_STR_LEN, "%s-%d-%d", prefix, loc_id, rmt_id);
  if(strlen(name) >= SHORT_STR_LEN - 1)p_die("prefix too long\n");
  p_info("publish qp :%s\n", name);
  mt_set(ctx->mi, name, (char *)&info, sizeof(struct qp_info));
}



void server_conn_init(struct server_ctx *ctx){
  int res, i;
  int server_num, client_num, id;
  server_num = ctx->param->server_num;
  client_num = ctx->param->client_num;
  id = ctx->id;

  //get the device and alloc the pd
  {
    struct ibv_device **dev_list;
    int num_dev;
    dev_list = ibv_get_device_list(&num_dev);
    assert(dev_list > 0);
    ctx->ctx = ibv_open_device(dev_list[0]);
    assert(ctx->ctx);
    res = ibv_query_port(ctx->ctx, ctx->param->ib_port, &ctx->port_attr);
    assert(!res);
    ctx->pd = ibv_alloc_pd(ctx->ctx);
    assert(ctx->pd);
  }
  //create oll the QP of servers
  {
    ctx->server_qp = (struct ibv_qp **)malloc(server_num * sizeof(struct ibv_qp *));
    ctx->server_send_cq = (struct ibv_cq **)malloc(server_num * sizeof(struct ibv_cq *));
    assert(ctx->server_qp && ctx->server_send_cq);
    ctx->server_recv_cq = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
    assert(ctx->server_recv_cq);
    for(i = 0; i < server_num; i++){
      if(i == id)continue;
      ctx->server_send_cq[i] = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
      assert(ctx->server_send_cq[i]);
      ctx->server_qp[i] = create_qp_init(ctx->pd, ctx->server_recv_cq, ctx->server_send_cq[i], ctx->param->ib_port);
      assert(ctx->server_qp[i]);
    }
  }
  //create all the QP of client
  {
    ctx->client_qp = (struct ibv_qp **)malloc(client_num * sizeof(struct ibv_qp *));
    ctx->client_send_cq = (struct ibv_cq **)malloc(client_num * sizeof(struct ibv_cq *));
    assert(ctx->client_qp && ctx->client_send_cq);
    ctx->client_recv_cq = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
    assert(ctx->client_recv_cq);
    for(i = 0; i < client_num; i++){
      ctx->client_send_cq[i] = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
      assert(ctx->client_send_cq[i]);
      ctx->client_qp[i] = create_qp_init(ctx->pd, ctx->client_recv_cq, 
                            ctx->client_send_cq[i], ctx->param->ib_port);
      assert(ctx->client_qp[i]);
    }
  }
  //publish all qps
  {
    //publish server QP
    for(i = 0; i < server_num; i++){
      if(i == id)continue;
      publish_qp_info(ctx, "server", ctx->server_qp[i], id, i);
    }
    //publish client QP
    for(i = 0; i < client_num; i++){
      publish_qp_info(ctx, "client", ctx->client_qp[i], id, i);
    }
  }
  //get all server qps
  {
    ctx->server_qp_info = (struct qp_info *)malloc(server_num * sizeof(struct qp_info));
    assert(ctx->server_qp_info);
    for(i = 0; i < server_num; i++){
      if(i == id)continue;
      get_qp_info(ctx->mi, "server", i, id, &ctx->server_qp_info[i]);
    }
  }
  //connect all server qps
  {
    for(i = 0; i < server_num; i++){
      if(i == id)continue;
      connect_qp(ctx->param->ib_port, ctx->server_qp[i], &ctx->server_qp_info[i]);
    }
  }
}

struct ibv_mr *reg_publish_mr(struct server_ctx *ctx, char *name, void *addr, long size){
  struct ibv_mr *mr;
  struct mr_info mr_info;
  assert(addr);
  mr = ibv_reg_mr(ctx->pd, addr, size, MR_PERMISSION);
  assert(mr);
  mr_info.size = size;
  mr_info.addr = (unsigned long)addr;
  mr_info.rkey = mr->rkey;
  mt_set(ctx->mi, name, (char *)&mr_info, sizeof(struct mr_info));
  return mr;
}

void sync_server_mr(struct server_ctx *ctx){
  char name[SHORT_STR_LEN];
  struct mr_info *inf;
  int i, server_num;

  server_num = ctx->param->server_num;
  snprintf(name, SHORT_STR_LEN, "smem-%d", ctx->id);
  ctx->pm_mr = reg_publish_mr(ctx, name, ctx->pm.addr, ctx->pm.size);
  ctx->server_mr_info = (struct mr_info* )malloc(server_num * sizeof(struct mr_info));
  for(i = 0; i < server_num; i++){
    snprintf(name, SHORT_STR_LEN, "smem-%d", i);
    while(mt_get(ctx->mi, name, (char **)&inf, NULL)) usleep(1000);
    memcpy(&ctx->server_mr_info[i], inf, sizeof(struct mr_info));
    p_info("get mr info %s addr :%lx len: %ld\n", name, 
            ctx->server_mr_info[i].addr, ctx->server_mr_info[i].size);
    free(inf);
  }
}

static void set_server_chash(struct mc_ctx *mi, int id,unsigned long h){
  char name[SHORT_STR_LEN];
  snprintf(name, SHORT_STR_LEN, "server-hash-%d", id);
  mt_set(mi, name, (char *)&h, sizeof(h));
}
static unsigned long get_server_chash(struct mc_ctx *mi, int id){
  char name[SHORT_STR_LEN];
  unsigned long *ph, h;
  snprintf(name, SHORT_STR_LEN, "server-hash-%d", id);
  while(mt_get(mi, name, (char **)&ph, NULL))usleep(1000);
  h = *ph;
  free(ph);
  return h;
}
void sync_server_hash(struct server_ctx *ctx){
  int server_num = ctx->param->server_num, i;
  set_server_chash(ctx->mi, ctx->id, ctx->super->hash);
  ctx->chash = (unsigned long *)malloc(server_num * sizeof(unsigned long));
  //p_info("%lx\n",ctx->super->hash);
  for(i = 0; i < server_num ;i++){
    ctx->chash[i] = get_server_chash(ctx->mi, i);
    //p_info("%lx\n", ctx->chash[i]);
  }
}
